/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package cloud.coder.remote.grpc.client.service;

import cloud.coder.remote.boot.BootService;
import cloud.coder.remote.boot.DefaultImplementor;
import cloud.coder.remote.boot.DefaultNamedThreadFactory;
import cloud.coder.remote.boot.ServiceManager;
import cloud.coder.remote.config.Config;
import cloud.coder.remote.grpc.client.GRPCChannelListener;
import cloud.coder.remote.grpc.client.GRPCChannelManager;
import cloud.coder.remote.grpc.client.GRPCChannelStatus;
import cloud.coder.remote.os.OSUtil;
import cloud.coder.remote.util.RunnableWithExceptionProtection;
import cloud.coder.remote.util.StringUtil;
import io.grpc.Channel;
import io.grpc.examples.network.common.Commands;
import io.grpc.examples.network.common.KeyStringValuePair;
import io.grpc.examples.network.management.InstancePingPkg;
import io.grpc.examples.network.management.InstanceProperties;
import io.grpc.examples.network.management.ManagementServiceGrpc;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

@DefaultImplementor
public class ServiceManagementClient implements BootService, Runnable, GRPCChannelListener {
    private static final Logger LOGGER = LoggerFactory.getLogger(ServiceManagementClient.class);
    private long reportInterval = 60L;
    private static List<KeyStringValuePair> SERVICE_INSTANCE_PROPERTIES;

    private volatile GRPCChannelStatus status = GRPCChannelStatus.DISCONNECT;
    private volatile ManagementServiceGrpc.ManagementServiceBlockingStub managementServiceBlockingStub;
    private volatile ScheduledFuture<?> heartbeatFuture;
    private volatile AtomicInteger sendPropertiesCounter = new AtomicInteger(0);

    @Override
    public void statusChanged(GRPCChannelStatus status) {
        if (GRPCChannelStatus.CONNECTED.equals(status)) {
            Channel channel = ServiceManager.INSTANCE.findService(GRPCChannelManager.class).getChannel();
            managementServiceBlockingStub = ManagementServiceGrpc.newBlockingStub(channel);
        } else {
            managementServiceBlockingStub = null;
        }
        this.status = status;
    }

    @Override
    public void prepare() {
        ServiceManager.INSTANCE.findService(GRPCChannelManager.class).addChannelListener(this);

        SERVICE_INSTANCE_PROPERTIES = new ArrayList<>();

        for (String key : Config.INSTANCE_PROPERTIES.keySet()) {
            SERVICE_INSTANCE_PROPERTIES.add(KeyStringValuePair.newBuilder()
                    .setKey(key)
                    .setValue(Config.INSTANCE_PROPERTIES.get(key))
                    .build());
        }

        Config.INSTANCE_NAME = StringUtil.isEmpty(Config.INSTANCE_NAME)
                ? UUID.randomUUID().toString().replaceAll("-", "") + "@" + OSUtil.getIPV4()
                : Config.INSTANCE_NAME;

        reportInterval = Config.HEARTBEAT_PERIOD * Config.PROPERTIES_REPORT_PERIOD_FACTOR;
    }

    @Override
    public void boot() {
        heartbeatFuture = new ScheduledThreadPoolExecutor(1,
                new DefaultNamedThreadFactory("ServiceManagementClient")
        ).scheduleAtFixedRate(
                new RunnableWithExceptionProtection(
                        this,
                        t -> LOGGER.error("unexpected exception.", t)
                ), 0, Config.HEARTBEAT_PERIOD,
                TimeUnit.SECONDS
        );
    }

    @Override
    public void onComplete() {
    }

    @Override
    public void shutdown() {
        heartbeatFuture.cancel(true);
    }

    @Override
    public void run() {
        LOGGER.debug("ServiceManagementClient running, status:{}.", status);

        if (GRPCChannelStatus.CONNECTED.equals(status)) {
            try {
                if (managementServiceBlockingStub != null) {
                    if (Math.abs(sendPropertiesCounter.getAndAdd(1)) % reportInterval == 0) {

                        managementServiceBlockingStub
                                .withDeadlineAfter(Config.GRPC_UPSTREAM_TIMEOUT, TimeUnit.SECONDS)
                                .reportInstanceProperties(InstanceProperties.newBuilder()
                                        .setService(Config.SERVICE_NAME)
                                        .setServiceInstance(Config.INSTANCE_NAME)
                                        .addAllProperties(OSUtil.buildOSInfo(
                                                Config.IPV4_LIST_SIZE))
                                        .addAllProperties(SERVICE_INSTANCE_PROPERTIES)
                                        .build());
                    } else {
                        final Commands commands = managementServiceBlockingStub.withDeadlineAfter(
                                Config.GRPC_UPSTREAM_TIMEOUT, TimeUnit.SECONDS
                        ).keepAlive(InstancePingPkg.newBuilder()
                                .setService(Config.SERVICE_NAME)
                                .setServiceInstance(Config.INSTANCE_NAME)
                                .build());

                        LOGGER.info(commands.toString());
                    }
                }
            } catch (Throwable t) {
                LOGGER.error("ServiceManagementClient execute fail." + t.getMessage(), t);
                ServiceManager.INSTANCE.findService(GRPCChannelManager.class).reportError(t);
            }
        }
    }
}
